# transformer_weather_forecast.py

import torch
import torch.nn as nn
import pandas as pd
from sklearn.preprocessing import MinMaxScaler

import numpy as np


# 生成模拟天气数据
def generate_weather_data(num_days):
    np.random.seed(0)
    dates = pd.date_range(start='2023-01-01', periods=num_days)
    temperatures = np.random.normal(loc=20, scale=5, size=num_days)  # 平均20度，标准差5
    humidities = np.random.uniform(low=30, high=90, size=num_days)  # 湿度范围30%到90%
    wind_speeds = np.random.uniform(low=0, high=15, size=num_days)  # 风速范围0到15 m/s

    weather_data = pd.DataFrame({
        'date': dates,
        'temperature': temperatures,
        'humidity': humidities,
        'wind_speed': wind_speeds
    })

    return weather_data


class WeatherTransformer(nn.Module):
    def __init__(self, input_dim, model_dim, num_heads, num_layers):
        super(WeatherTransformer, self).__init__()
        self.embedding = nn.Linear(input_dim, model_dim)
        self.transformer = nn.Transformer(model_dim, num_heads, num_layers)
        self.fc = nn.Linear(model_dim, input_dim)  # 输出维度与输入维度相同

    def forward(self, x, tgt):
        x = self.embedding(x)
        tgt = self.embedding(tgt)
        x = self.transformer(x, tgt)
        x = self.fc(x)
        return x


# 数据预处理
def preprocess_data(weather_data):
    scaler = MinMaxScaler()
    scaled_data = scaler.fit_transform(weather_data[['temperature', 'humidity', 'wind_speed']])
    return scaled_data, scaler


# 训练模型
def train_model(model, data, epochs=1000):
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    for epoch in range(epochs):
        model.train()
        optimizer.zero_grad()
        output = model(data, data)
        loss = criterion(output, data)
        loss.backward()
        optimizer.step()
        if epoch % 10 == 0:
            print(f'Epoch {epoch}, Loss: {loss.item()}')


# 主程序
if __name__ == "__main__":
    # 生成天气数据
    weather_data = generate_weather_data(150)
    scaled_data, scaler = preprocess_data(weather_data)

    # 转换为 PyTorch 张量
    input_tensor = torch.tensor(scaled_data[:-1], dtype=torch.float32).unsqueeze(
        1)  # 形状调整为 (seq_len, batch_size, input_dim)

    # 初始化模型
    model = WeatherTransformer(input_dim=3, model_dim=64, num_heads=4, num_layers=2)

    # 训练模型
    train_model(model, input_tensor)

    # 预测未来天气
    model.eval()
    with torch.no_grad():
        future_weather = model(input_tensor[-1].unsqueeze(0), input_tensor[-1].unsqueeze(0))  # 使用最后一天的数据进行预测
        future_weather = scaler.inverse_transform(future_weather.numpy().reshape(-1, 3))  # 反归一化
        print("未来天气预测（温度, 湿度, 风速）:", future_weather)
